import torch
import torch.nn as nn
import torch.nn.functional as F


class ActorContinue(nn.Module):
    def __init__(self, input_size, action_size, action_bound, args):
        super(ActorContinue, self).__init__()
        self.layer1 = nn.Linear(input_size, args.num_units_1)
        self.layer1.weight.data.normal_(0, 1)
        self.layer2 = nn.Linear(args.num_units_1, args.num_units_2)
        self.layer2.weight.data.normal_(0, 1)
        self.layer_out = nn.Linear(args.num_units_2, action_size)
        self.layer_out.weight.data.normal_(0, 1)

        self.action_bound = action_bound
    
    def forward(self, actor_input):
        x = F.relu(self.layer1(actor_input))
        x = F.relu(self.layer2(x))
        x = self.layer_out(x)
        x = torch.tanh(x) * self.action_bound
        return x

# test
# from common.arguments import get_common_args
# args = get_common_args()
# act = ActorContinue(10, 1, 2, args)
# print(act.state_dict())
